import torch
import triton
import triton.language as tl

''' eval result:
🚀 Speedup=1.09x size 16x16x16 triton_time=0.0061 torch_time=0.0066
🚀 Speedup=1.05x size 32x32x32 triton_time=0.0062 torch_time=0.0065
🚀 Speedup=1.07x size 64x64x64 triton_time=0.0063 torch_time=0.0068
🚀 Speedup=1.08x size 128x128x128 triton_time=0.0064 torch_time=0.0069
🐌 Slowdown=0.49x size 512x2018x512 triton_time=0.0381 torch_time=0.0186
🐌 Slowdown=0.52x size 1024x4096x4096 triton_time=0.0969 torch_time=0.0505
'''

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

    offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    offs_k = tl.arange(0, BLOCK_SIZE)
    
    a_ptrs = a_ptr + offs_m[:, None] * stride_ak + offs_k[None, :]
    b_ptrs = b_ptr + offs_k[:, None] * stride_bn + offs_n[None, :]
    
    accumulator = tl.zeros((BLOCK_SIZE, BLOCK_SIZE), dtype=tl.float32)
    
    for k in range(0, tl.cdiv(K, BLOCK_SIZE)):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < (K - k * BLOCK_SIZE), other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < (K - k * BLOCK_SIZE), other=0.0)
        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE
        b_ptrs += BLOCK_SIZE * stride_bn
    
    c_ptrs = c_ptr + offs_m[:, None] * stride_cn + offs_n[None, :]
    
    tl.store(c_ptrs, accumulator, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

def triton_matmul(a, b, c=None):
    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
    M, K = a.shape
    K, N = b.shape
    if c is None:
        c = torch.empty((M, N), device=a.device, dtype=torch.float32)
    
    grid = lambda meta: (triton.cdiv(M, meta['BLOCK_SIZE']), triton.cdiv(N, meta['BLOCK_SIZE']), )
    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        M, a.stride(0),
        K, b.stride(0),
        M, c.stride(0),
        BLOCK_SIZE=64
    )
    return c

# 使用 benchmark 比较 Triton 和 cuBLAS
def run_benchmark():
    from triton.testing import do_bench
    # 测试不同尺寸的矩阵
    sizes = [(16, 16, 16), (32, 32, 32), (64, 64, 64), (128, 128, 128), (512, 512, 2018), (1024, 4096, 4096)]

    # in_dtype = torch.float8_e4m3fn
    in_dtype = torch.float16
    out_dtype = torch.float16

    for (M, N, K) in sizes:
        a = torch.randn((M, K), device='cuda', dtype=torch.float32).to(in_dtype)
        b = torch.randn((K, N), device='cuda', dtype=torch.float32).to(in_dtype)
        c = torch.empty((M, N), device='cuda', dtype=torch.float32).to(out_dtype)

        def triton_fn():
            return triton_matmul(a, b, c)
        def torch_fn():
            return torch.matmul(a, b, out=c)
        
        triton_time = do_bench(triton_fn, warmup=100, rep=500)
        torch_time = do_bench(torch_fn, warmup=100, rep=500)
        speedup = 1.0/(triton_time/torch_time)
        if speedup > 1.0:
            print(f'🚀 Speedup={speedup:.2f}x size {M}x{K}x{N} {triton_time=:.4f} {torch_time=:.4f}')
        else:
            print(f'🐌 Slowdown={speedup:.2f}x size {M}x{K}x{N} {triton_time=:.4f} {torch_time=:.4f}')

def test_matmul():
    torch.manual_seed(0)
    
    # 测试不同尺寸的矩阵
    sizes = [(16, 16, 16), (32, 32, 32), (64, 64, 64), (128, 128, 128), (512, 512, 2048)]
    tolerances = [(16, 1e-5), (32, 1e-5), (64, 1e-4), (128, 1e-4), (512, 1e-3)]
    
    for (M, N, K), (_, tol) in zip(sizes, tolerances):
        # # randn 无法通过测试，但算法是对的，随机值的误差就会比较大
        # a = torch.randn((M, K), device='cuda', dtype=torch.float32)
        # b = torch.randn((K, N), device='cuda', dtype=torch.float32)

        # 可以通过
        a = torch.full((M, K), 0.01, device='cuda', dtype=torch.float32)
        b = torch.full((K, N), 0.02, device='cuda', dtype=torch.float32)

        # 计算 Triton 结果
        triton_output = triton_matmul(a, b)
        
        # 计算 PyTorch 参考结果
        torch_output = torch.matmul(a, b)

        # print(f'{a=} \n {b=} \n {triton_output=} \n {torch_output=}')

        # 验证结果
        if torch.allclose(triton_output, torch_output, rtol=tol, atol=tol):
            print(f"✅ Test passed for size {M}x{K}x{N} (tolerance={tol})")
        else:
            max_diff = torch.max(torch.abs(triton_output - torch_output)).item()
            print(f"❌ Test failed for size {M}x{K}x{N}: Max diff={max_diff:.6f}")

if __name__ == "__main__":
    # 检查 GPU 可用性
    if not torch.cuda.is_available():
        print("CUDA not available. Skipping tests.")
    else:
        print("Starting matrix multiplication tests...")
        test_matmul()
        run_benchmark()
